from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import torch
from bertviz import head_view,model_view 
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
)
from pathlib import Path
import os

class SentimentAnalyzer:
    """封装好的情感分析工具"""
    def __init__(self, model_path):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
        self.device = 0 if torch.cuda.is_available() else 'cpu'
        self.model.to(self.device)
        
    def predict(self, text):
        tokenizer = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
        tokens = self.tokenizer.convert_ids_to_tokens(tokenizer["input_ids"][0])
        
        inputs = {k: v.to(self.device) for k, v in tokenizer.items()}
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            attentions = outputs.attentions

        
        probs = torch.softmax(outputs.logits, dim=-1)[0]
        pred_label = torch.argmax(probs).item()
        
        return {
            "text": text,
            "label": self.model.config.id2label[pred_label],
            "confidence": probs[pred_label].item(),
            "probabilities": {
                self.model.config.id2label[i]: prob.item() 
                for i, prob in enumerate(probs)
            },
            "attentions" : attentions,
            'tokenizer' : tokenizer,
            'tokens' : tokens
        }


BASE_DIR = str(Path(__file__).parent)

app = FastAPI(root_path=BASE_DIR)

app.mount("/static", StaticFiles(directory=BASE_DIR), name="static")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

sa = SentimentAnalyzer(os.path.join(BASE_DIR,'result'))

class TextRequest(BaseModel):
    text: str


@app.get("/attention")
async def get_attention_html():
    """返回注意力可视化的HTML文件"""
    
    file_path = BASE_DIR + "\\attention.html"
    if not file_path.exists():
        return {"error": "Attention visualization HTML file not found."}
    
    return FileResponse(file_path)


@app.post("/text_predict")
async def predict_text(request: TextRequest):
    """文本情感分析API"""
    res = sa.predict(request.text)
    result = res['label']
    html = head_view(attention=res['attentions'], tokens=res['tokens'], html_action='return')
    with open(os.path.join(BASE_DIR, 'attention.html'), 'w') as f:
        f.write(html.data)

    return {'predicted_class': result, 'vis_data': html.data}
